import shutil
import os
import re

import numpy as np

from tqdm import tqdm
from datetime import datetime


def synchronize_videoandradar_folders(allvideofolder='./VIDEO/', 
                                      allradardatafolder='./RADAR/',
                                      videodt_format = "%Y%m%d%H%M",
                                      radardt_format = "%Y%m%d%H%M%S"):
    """
    function: Put video files and radar files with the same capture time into the same folder
    args: allvideofolder: the root path of video data
          allradardatafolder: the root path of radar data
          videodt_format: the tine format of video
          radardt_format: the time format of radar
        
    return: put two into one filefolder
    
    Example usage:
    synchronize_folders('./VIDEO/', './RADAR/')
    """
    # List video and radar folders
    allvideofolder_list = os.listdir(allvideofolder)
    allvideofolder_list = sorted(allvideofolder_list, key=lambda x: int(re.findall(r'\d+', x)[0]))

    allradardatafolder_list = os.listdir(allradardatafolder)
    allradardatafolder_list = sorted(allradardatafolder_list, key=lambda x: int(re.findall(r'\d+', x)[0]))

    # Find matching folders
    same_folder = []
    for i in range(len(allvideofolder_list)):
        video_strat_timedt = datetime.strptime(str(allvideofolder_list[i]), videodt_format)
        video_strat_timestamp = video_strat_timedt.timestamp()
        
        for j in range(len(allradardatafolder_list)):
            radar_strat_timedt = datetime.strptime(str(allradardatafolder_list[j]), radardt_format)
            radar_strat_timestamp = radar_strat_timedt.timestamp()

            if abs(video_strat_timestamp - radar_strat_timestamp) <= 100:
                same_folder.append([allvideofolder_list[i], allradardatafolder_list[j]])

    newpath = './SYNCHRONIZATION'

    for i in range(len(same_folder)):
        videofolderpath = os.path.join(allvideofolder, same_folder[i][0])
        radarfolderpath = os.path.join(allradardatafolder, same_folder[i][1])
        
        bothsensorspath = os.path.join(newpath, same_folder[i][0])
        newvideofolderpath = os.path.join(newpath, same_folder[i][0], 'VIDEO')
        newradarfolderpath = os.path.join(newpath, same_folder[i][0], 'RADAR')
        
        if not os.path.exists(bothsensorspath):
            os.mkdir(bothsensorspath)
            
            if not os.path.exists(newvideofolderpath):
                os.mkdir(newvideofolderpath)
                shutil.move(videofolderpath, newvideofolderpath)
            
            if not os.path.exists(newradarfolderpath):
                os.mkdir(newradarfolderpath)
                shutil.move(radarfolderpath, newradarfolderpath)



def synchronize_radar_times(SYNCHRO_FOLDER='./TRANSFER_STATION_SYNCHRONIZATION/'):

    """
    function: make video file and radar file time synchronize
    args: SYNCHRO_FOLDER: the root path of video data and radar data

    return: SAVE_TIME = os.path.join(SYNCHRO_FOLDER, 
                                 SYNCHRO_FOLDER_LIST[i],
                                'initial_sync_index.npy')
            make such a file to save same timestamp

    Example usage:
    synchronize_times('./TRANSFER_STATION_SYNCHRONIZATION/')
    """

    SYNCHRO_FOLDER_LIST = os.listdir(SYNCHRO_FOLDER)
    SYNCHRO_FOLDER_LIST = sorted(SYNCHRO_FOLDER_LIST, 
                                 key=lambda x: int(re.findall(r'\d+', 
                                                              x)[0]))

    for i in range(len(SYNCHRO_FOLDER_LIST)):
        RADAR_TIME = np.load(os.path.join(SYNCHRO_FOLDER,
                                          SYNCHRO_FOLDER_LIST[i],
                                          'RADAR',
                                          os.listdir(os.path.join(SYNCHRO_FOLDER, 
                                                                  SYNCHRO_FOLDER_LIST[i], 
                                                                  'RADAR'))[0],
                                          os.listdir(os.path.join(SYNCHRO_FOLDER, 
                                                                  SYNCHRO_FOLDER_LIST[i], 
                                                                  'RADAR'))[0] + '_frame.npy'))

        VIDEO_TIME = np.load(os.path.join(SYNCHRO_FOLDER, 
                                          SYNCHRO_FOLDER_LIST[i],
                                          'VIDEO',
                                          SYNCHRO_FOLDER_LIST[i],
                                          'frame.pkl'), 
                                          allow_pickle=True)
        VIDEO_TIME = np.vectorize(lambda x: x.timestamp())(VIDEO_TIME)

        label_timestamp = np.load(os.path.join(SYNCHRO_FOLDER, 
                                               SYNCHRO_FOLDER_LIST[i],
                                               'VIDEO',
                                               SYNCHRO_FOLDER_LIST[i],
                                               'timestamp.npy'))

        SAVE_TIME = os.path.join(SYNCHRO_FOLDER, 
                                 SYNCHRO_FOLDER_LIST[i],
                                'initial_sync_index.npy')

        sameindexNP = []
        VIDEO_TIME = np.array(VIDEO_TIME)
        RADAR_TIME = np.array(RADAR_TIME)

        video_time_round = np.round(VIDEO_TIME, 1)
        radar_time_round = np.round(RADAR_TIME, 1)

        for i in range(len(VIDEO_TIME)):
            index = np.where(video_time_round[i] == radar_time_round)[0]
            if len(index) > 0:
                j = index[0]
                sameindexNP.append([i, VIDEO_TIME[i], j, RADAR_TIME[j]])

        sameindexNP = np.array(sameindexNP)

        mid_time_index = []
        for i in range(len(label_timestamp)):
            start_time = label_timestamp[i][0]
            end_time = label_timestamp[i][1]
            one = []
            for j in range(len(sameindexNP)):
                video_time_index = sameindexNP[j][0]
                video_time = sameindexNP[j][1]
                radar_time_index = sameindexNP[j][2]
                if start_time <= video_time <= end_time:
                    one.append([video_time_index, radar_time_index])

            one = np.array(one)
            mid_time_index.append(one)

        mid_time_index = [x for x in mid_time_index if x.any()]

        if not os.path.exists(SAVE_TIME):
            np.save(SAVE_TIME, mid_time_index, allow_pickle=True)



